import sys
sys.path.append("core_trig_dry_65")
from richards_equation_1D import volMoistureAllNodes
from parameters_mpc import *
import numpy as np
from matplotlib import pyplot as plt
from model_simulation import simulate_Richards_equation
from crop_coefficient_calculation import generate_Kc
totalDepth, axialNodes, totalNodes = spatialVariables_1D_act()

# Load the  soil parameters for the MZs
soilPars_mz1 = ManagementZone_1()
soilPars_mz2 = ManagementZone_2()
soilPars_mz3 = ManagementZone_3()

# Relevant for the simulation
sequenceLength = 5  # Should not be changed
horizonLength_orig = 14  # Can be tuned
horizonLength_weather = int(sequenceLength + horizonLength_orig)

# Load the weather conditions
_weather_conditions = np.loadtxt("./relevant_data/Weather Data_dry_n.txt")
_average_temps = _weather_conditions[:, 0]
_crop_coeff = generate_Kc(_average_temps)
_rain = _weather_conditions[:, 1]
_ref_Evap = _weather_conditions[:, 2]


# _predicted_rain = -1*np.loadtxt("Predicted_rain_a.txt")
# _predicted_rain = (_predicted_rain)/1000
# for i in range(len(_ref_Evap)):
#     if _ref_Evap[i]< 1.04:
#         _ref_Evap[i] = 1.04
#     if _ref_Evap[i] > 9.0:
#         _ref_Evap[i] = 9.0
#     pass

# _crop_coeff[0:20]=0.90


def compute_root_zone_moisture(Xk, rooting_depth):

    if rooting_depth == 1.00:
        root_zone_head = 0.10*((1/6)*(Xk[0] + Xk[1] + Xk[2] + Xk[3] + Xk[4] + Xk[5])) + \
            0.20*((1/6)*(Xk[5] + Xk[6] + Xk[7] + Xk[8] + Xk[9] + Xk[10])) + \
            0.30*((1/11)*(Xk[10] + Xk[11] + Xk[12] + Xk[13] + Xk[14] + Xk[15] + Xk[16] + Xk[17] + Xk[18] + Xk[19] + Xk[20])) + \
            0.40*((1/11)*(Xk[20] + Xk[21] + Xk[22] + Xk[23] + Xk[24] +
                  Xk[25] + Xk[26] + Xk[27] + Xk[28] + Xk[29] + Xk[30]))
    else:
        root_zone_head = 0.10*((1/6)*(Xk[10] + Xk[11] + Xk[12] + Xk[13] + Xk[14] + Xk[15])) +\
            0.20*((1/6)*(Xk[15] + Xk[16] + Xk[17] + Xk[18] + Xk[19] + Xk[20])) +\
            0.30*((1/6)*(Xk[20] + Xk[21] + Xk[22] + Xk[23] + Xk[24] + Xk[25])) + \
            0.40*((1/6)*(Xk[25] + Xk[26] + Xk[27] + Xk[28] + Xk[29] + Xk[30]))

    return root_zone_head

# Converts the root zone pressure head to soil moisture


def obtain_soil_moisture(head, soilPars):
    vol_moist = volMoistureAllNodes(head, soilPars)
    return vol_moist


# Load evapotranspiration and crop coefficient
referenceEvap_all_mz1 = _ref_Evap/(86400*1000)
referenceEvap_all_act_mz1 = _ref_Evap/(86400*1000)
cropCoefficient_all_mz1 = _crop_coeff

referenceEvap_all_mz2 = _ref_Evap/(86400*1000)
referenceEvap_all_act_mz2 = _ref_Evap/(86400*1000)
cropCoefficient_all_mz2 = _crop_coeff

referenceEvap_all_mz3 = _ref_Evap/(86400*1000)
referenceEvap_all_act_mz3 = _ref_Evap/(86400*1000)
cropCoefficient_all_mz3 = _crop_coeff
rain_all = -1*(_rain/1000)

#predicted_rain = rain_all

predicted_rain = np.loadtxt("./relevant_data/Predicted_rain_n.txt")

simulationPeriod = 123  # Can be tuned
rooting_depths = np.zeros(160)
rooting_depths[0:72] = 0.50
rooting_depths[72:] = 1.00

lai_factors_all = np.loadtxt("./relevant_data/LAI_factors.txt")


def compute_irrigation_rate(soilMoist_current, soilMoist_fc, irrigation_decision, rootZone_depth):
    if irrigation_decision == 0:
        return 0
    else:
        return (soilMoist_fc - soilMoist_current) * rootZone_depth


def compute_irrigation_rate_mz1(soilMoist_current, soilMoist_fc, irrigation_decision, rootZone_depth):
    if irrigation_decision == 0:
        return 0
    else:
        irrigRate = (soilMoist_fc - soilMoist_current) * rootZone_depth
        if irrigRate > 0.0518:
            irrigRate_actual = 0.0518
        else:
            irrigRate_actual = irrigRate
        return irrigRate_actual


def compute_irrigation_rate_mz2(soilMoist_current, soilMoist_fc, irrigation_decision, rootZone_depth):
    if irrigation_decision == 0:
        return 0
    else:
        irrigRate = (soilMoist_fc - soilMoist_current) * rootZone_depth
        if irrigRate > 0.0596:
            irrigRate_actual = 0.0596
        else:
            irrigRate_actual = irrigRate
        return irrigRate_actual


def compute_irrigation_rate_mz3(soilMoist_current, soilMoist_fc, irrigation_decision, rootZone_depth):
    if irrigation_decision == 0:
        return 0
    else:
        irrigRate = (soilMoist_fc - soilMoist_current) * rootZone_depth
        if irrigRate > 0.0622:
            irrigRate_actual = 0.0622
        else:
            irrigRate_actual = irrigRate
        return irrigRate_actual


# def compute_irrigation_rate_mz1(soilMoist_current, soilMoist_fc, irrigation_decision, rootZone_depth):
#     if irrigation_decision == 0:
#         return 0
#     else:
#         irrigRate = (soilMoist_fc - soilMoist_current) * rootZone_depth
#         if irrigRate > 0.030:
#             irrigRate_actual = 0.030
#         else:
#             irrigRate_actual = irrigRate
#         return irrigRate_actual

# def compute_irrigation_rate_mz2(soilMoist_current, soilMoist_fc, irrigation_decision, rootZone_depth):
#     if irrigation_decision == 0:
#         return 0
#     else:
#         irrigRate = (soilMoist_fc - soilMoist_current) * rootZone_depth
#         if irrigRate > 0.030:
#             irrigRate_actual = 0.0300
#         else:
#             irrigRate_actual = irrigRate
#         return irrigRate_actual

# def compute_irrigation_rate_mz3(soilMoist_current, soilMoist_fc, irrigation_decision, rootZone_depth):
#     if irrigation_decision == 0:
#         return 0
#     else:
#         irrigRate = (soilMoist_fc - soilMoist_current) * rootZone_depth
#         if irrigRate > 0.0340:
#             irrigRate_actual = 0.0340
#         else:
#             irrigRate_actual = irrigRate
#         return irrigRate_actual


def compute_irrigation_decision(soilMoist_current, soilMoist_thresh):
    if soilMoist_current <= soilMoist_thresh:
        return 1
    else:
        return 0


def determine_irrigation_rates(irrig_decision_mz1, irrig_decision_mz2, irrig_decision_mz3, current_rootzone_moist_mz1,
                               UpperZone_mz1, current_rootzone_moist_mz2, UpperZone_mz2, current_rootzone_moist_mz3, UpperZone_mz3, rd_currentSE_mz1):

    # Determine the binding irrigation decision

    if irrig_decision_mz1+irrig_decision_mz2+irrig_decision_mz3 == 0:
        binding_decision = 0
        current_irrigationrate_mz1 = compute_irrigation_rate_mz1(
            current_rootzone_moist_mz1, UpperZone_mz1, binding_decision, rd_currentSE_mz1)
        current_irrigationrate_mz2 = compute_irrigation_rate_mz2(
            current_rootzone_moist_mz2, UpperZone_mz2, binding_decision, rd_currentSE_mz1)
        current_irrigationrate_mz3 = compute_irrigation_rate_mz3(
            current_rootzone_moist_mz3, UpperZone_mz3, binding_decision, rd_currentSE_mz1)
        return current_irrigationrate_mz1, current_irrigationrate_mz2, current_irrigationrate_mz3, binding_decision

    else:
        binding_decision = 1  # If at least one of the irrigation decisions is 1
        current_irrigationrate_mz1 = compute_irrigation_rate_mz1(
            current_rootzone_moist_mz1, UpperZone_mz1, binding_decision, rd_currentSE_mz1)
        current_irrigationrate_mz2 = compute_irrigation_rate_mz2(
            current_rootzone_moist_mz2, UpperZone_mz2, binding_decision, rd_currentSE_mz1)
        current_irrigationrate_mz3 = compute_irrigation_rate_mz3(
            current_rootzone_moist_mz3, UpperZone_mz3, binding_decision, rd_currentSE_mz1)
        return current_irrigationrate_mz1, current_irrigationrate_mz2, current_irrigationrate_mz3, binding_decision


# Load the initial pressure head
allHeadValues_mz1 = np.loadtxt(
    "./relevant_data/Init_state_all_mz1_vrd_lai.txt")
allHeadValues_mz2 = np.loadtxt(
    "./relevant_data/Init_state_all_mz2_vrd_lai.txt")
allHeadValues_mz3 = np.loadtxt(
    "./relevant_data/Init_state_all_mz3_vrd_lai.txt")
simulationPeriod = 123
rootingDepth_init = 0.50

# Lists to store the simulation results
prescribedIrrigation_closedLoop_mz1 = []
prescribedIrrigationDecision_closedLoop_mz1 = []
stateTrajectory_closedLoop_mz1 = []
# root_zone_head_mz1 = compute_root_zone_moisture(allHeadValues_mz1)
stateTrajectory_closedLoop_mz1.append(compute_root_zone_moisture(
    obtain_soil_moisture(allHeadValues_mz1, soilPars_mz1), rootingDepth_init))

prescribedIrrigation_closedLoop_mz2 = []
stateTrajectory_closedLoop_mz2 = []
prescribedIrrigationDecision_closedLoop_mz2 = []
# root_zone_head_mz2 = compute_root_zone_moisture(allHeadValues_mz2)
stateTrajectory_closedLoop_mz2.append(compute_root_zone_moisture(
    obtain_soil_moisture(allHeadValues_mz2, soilPars_mz2), rootingDepth_init))

prescribedIrrigation_closedLoop_mz3 = []
stateTrajectory_closedLoop_mz3 = []
prescribedIrrigationDecision_closedLoop_mz3 = []
# root_zone_head_mz3 = compute_root_zone_moisture(allHeadValues_mz3)
stateTrajectory_closedLoop_mz3.append(compute_root_zone_moisture(
    obtain_soil_moisture(allHeadValues_mz3, soilPars_mz3), rootingDepth_init))

irrigationDecisions_closedLoop = []
# For the simulation of the actual plant, the Richards Equation in this work
allStates_mz1 = np.zeros((simulationPeriod+1, totalNodes))
allStates_mz1[0, :] = allHeadValues_mz1
allStates_mz2 = np.zeros((simulationPeriod+1, totalNodes))
allStates_mz2[0, :] = allHeadValues_mz2
allStates_mz3 = np.zeros((simulationPeriod+1, totalNodes))
allStates_mz3[0, :] = allHeadValues_mz3


# # Lower and upper bounds of the target zone for each MZ
# LowerZone_mz1 = 0.192
# UpperZone_mz1 = 0.280 #0.280
# LowerZone_mz2 = 0.192
# UpperZone_mz2 = 0.280 #0.280
# LowerZone_mz3 = 0.223
# UpperZone_mz3 = 0.300 #0.300

# Lower and upper bounds of the target zone for each MZ
LowerZone_mz1 = 0.176
UpperZone_mz1 = 0.280  # 0.280
LowerZone_mz2 = 0.176
UpperZone_mz2 = 0.280  # 0.280
LowerZone_mz3 = 0.209
UpperZone_mz3 = 0.300  # 0.300


for k in range(simulationPeriod):
    print("\n\n")
    print("------------Open-loop simulation for day", k+1)
    kc_init_mz1 = cropCoefficient_all_mz1[k:k+horizonLength_weather]
    et_init_mz1 = referenceEvap_all_mz1[k:k+horizonLength_weather]
    kc_init_mz2 = cropCoefficient_all_mz2[k:k+horizonLength_weather]
    et_init_mz2 = referenceEvap_all_mz2[k:k+horizonLength_weather]
    kc_init_mz3 = cropCoefficient_all_mz3[k:k+horizonLength_weather]
    et_init_mz3 = referenceEvap_all_mz3[k:k+horizonLength_weather]

    rain = rain_all[k: k+horizonLength_weather]
    #rain_predicted = predicted_rain[k : k+horizonLength_weather]
    rain_predicted = predicted_rain[k: k+horizonLength_weather]
    rain_predicted = predicted_rain[k, :]
    rooting_depth = rooting_depths[k:k+horizonLength_weather]
    lai_factors = lai_factors_all[k:k+horizonLength_weather]
    rd_currentSE_mz1 = rooting_depth[4]
    lai_currentSE_mz1 = lai_factors[4]

    current_rootzone_moist_mz1 = stateTrajectory_closedLoop_mz1[k]
    current_rootzone_moist_mz2 = stateTrajectory_closedLoop_mz2[k]
    current_rootzone_moist_mz3 = stateTrajectory_closedLoop_mz3[k]

    current_decision_mz1 = compute_irrigation_decision(
        current_rootzone_moist_mz1, LowerZone_mz1)
    current_decision_mz2 = compute_irrigation_decision(
        current_rootzone_moist_mz2, LowerZone_mz2)
    current_decision_mz3 = compute_irrigation_decision(
        current_rootzone_moist_mz3, LowerZone_mz3)

    current_irrigationrate_mz1, current_irrigationrate_mz2, current_irrigationrate_mz3, binding_decision = determine_irrigation_rates(current_decision_mz1, current_decision_mz2, current_decision_mz3,
                                                                                                                                      current_rootzone_moist_mz1, UpperZone_mz1, current_rootzone_moist_mz2, UpperZone_mz2, current_rootzone_moist_mz3, UpperZone_mz3, rd_currentSE_mz1)

    # prescribedIrrigationDecision_closedLoop_mz1.append(current_decision_mz1)
    # prescribedIrrigationDecision_closedLoop_mz2.append(current_decision_mz2)
    # prescribedIrrigationDecision_closedLoop_mz3.append(current_decision_mz3)

    if rain[4]!=0:
        binding_decision = 0
        current_irrigationrate_mz1 = current_irrigationrate_mz2 = current_irrigationrate_mz3 = 0
        pass

    prescribedIrrigationDecision_closedLoop_mz1.append(binding_decision)
    prescribedIrrigationDecision_closedLoop_mz2.append(binding_decision)
    prescribedIrrigationDecision_closedLoop_mz3.append(binding_decision)

    # current_irrigationrate_mz1 = compute_irrigation_rate_mz1(current_rootzone_moist_mz1, UpperZone_mz1, current_decision_mz1, rd_currentSE_mz1)
    # current_irrigationrate_mz2 = compute_irrigation_rate_mz2(current_rootzone_moist_mz2, UpperZone_mz2, current_decision_mz2, rd_currentSE_mz1)
    # current_irrigationrate_mz3 = compute_irrigation_rate_mz3(current_rootzone_moist_mz3, UpperZone_mz3, current_decision_mz3, rd_currentSE_mz1)

    # current_irrigationrate_mz1 = compute_irrigation_rate(current_rootzone_moist_mz1, UpperZone_mz1, current_decision_mz1, rd_currentSE_mz1)
    # current_irrigationrate_mz2 = compute_irrigation_rate(current_rootzone_moist_mz2, UpperZone_mz2, current_decision_mz2, rd_currentSE_mz1)
    # current_irrigationrate_mz3 = compute_irrigation_rate(current_rootzone_moist_mz3, UpperZone_mz3, current_decision_mz3, rd_currentSE_mz1)

    # if current_decision_mz1 == 1:
    #     current_irrigationrate_mz1 = current_irrigationrate_mz1 + 1*sum(rain_predicted[5:5+2])
    #     if current_irrigationrate_mz1 < 0:
    #         current_irrigationrate_mz1 = 0
    # else:
    #     current_irrigationrate_mz1 = current_irrigationrate_mz1

    # if current_decision_mz2 == 1:
    #     current_irrigationrate_mz2 = current_irrigationrate_mz2 + 1*sum(rain_predicted[5:5+2])
    #     if current_irrigationrate_mz2 < 0:
    #         current_irrigationrate_mz2 = 0
    # else:
    #     current_irrigationrate_mz2 = current_irrigationrate_mz2

    # if current_decision_mz3 == 1:
    #     current_irrigationrate_mz3 = current_irrigationrate_mz3 + 1*sum(rain_predicted[5:5+2])

    #     if current_irrigationrate_mz3 < 0:
    #         current_irrigationrate_mz3 = 0
    # else:
    #     current_irrigationrate_mz3 = current_irrigationrate_mz3

    # if current_irrigationrate_mz3 == 0:
    #     prescribedIrrigationDecision_closedLoop_mz3[k] = 0

    if binding_decision == 1:
        current_irrigationrate_mz1 = current_irrigationrate_mz1 + \
            1*sum(rain_predicted[4:4+4])
        if current_irrigationrate_mz1 < 0:
            current_irrigationrate_mz1 = 0
    else:
        current_irrigationrate_mz1 = current_irrigationrate_mz1

    if binding_decision == 1:
        current_irrigationrate_mz2 = current_irrigationrate_mz2 + \
            1*sum(rain_predicted[4:4+4])
        if current_irrigationrate_mz2 < 0:
            current_irrigationrate_mz2 = 0
    else:
        current_irrigationrate_mz2 = current_irrigationrate_mz2

    if binding_decision == 1:
        current_irrigationrate_mz3 = current_irrigationrate_mz3 + \
            1*sum(rain_predicted[4:4+4])

        if current_irrigationrate_mz3 < 0:
            current_irrigationrate_mz3 = 0
    else:
        current_irrigationrate_mz3 = current_irrigationrate_mz3

    # if current_irrigationrate_mz3 == 0:
    #     prescribedIrrigationDecision_closedLoop_mz3[k] = 0

    # current_irrigationrate_mz1 = current_irrigationrate_mz1 + sum(rain[4:4+4])
    # current_irrigationrate_mz2 = current_irrigationrate_mz2 + sum(rain[4:4+4])
    # current_irrigationrate_mz3 = current_irrigationrate_mz3 + sum(rain[4:4+4])
    # Without the rain
    prescribedIrrigation_closedLoop_mz1.append(current_irrigationrate_mz1)
    prescribedIrrigation_closedLoop_mz2.append(current_irrigationrate_mz2)
    prescribedIrrigation_closedLoop_mz3.append(current_irrigationrate_mz3)

    # Simulate to Richards equation
    x_currentSE_mz1 = allStates_mz1[k, :]
    u_currentSE_mz1 = (-current_irrigationrate_mz1 + rain[4])/86400
    kc_currentSE_mz1 = kc_init_mz1[4]
    et_currentSE_mz1 = et_init_mz1[4]

    currentStatesSE_mz1 = simulate_Richards_equation(
        x_currentSE_mz1, u_currentSE_mz1, et_currentSE_mz1, kc_currentSE_mz1, 4, rd_currentSE_mz1, lai_currentSE_mz1, soilPars_mz1)
    #rootzone_head_mz1 = compute_root_zone_moisture(currentStatesSE_mz1)
    rootzone_moist_mz1 = compute_root_zone_moisture(
        obtain_soil_moisture(currentStatesSE_mz1, soilPars_mz1), rd_currentSE_mz1)

    x_currentSE_mz2 = allStates_mz2[k, :]
    u_currentSE_mz2 = (-current_irrigationrate_mz2 + rain[4])/86400
    kc_currentSE_mz2 = kc_init_mz2[4]
    et_currentSE_mz2 = et_init_mz2[4]
    currentStatesSE_mz2 = simulate_Richards_equation(
        x_currentSE_mz2, u_currentSE_mz2, et_currentSE_mz2, kc_currentSE_mz2, 4, rd_currentSE_mz1, lai_currentSE_mz1, soilPars_mz2)
    #rootzone_head_mz2 = compute_root_zone_moisture(currentStatesSE_mz2)
    rootzone_moist_mz2 = compute_root_zone_moisture(
        obtain_soil_moisture(currentStatesSE_mz2, soilPars_mz2), rd_currentSE_mz1)

    x_currentSE_mz3 = allStates_mz3[k, :]
    u_currentSE_mz3 = (-current_irrigationrate_mz3 + rain[4])/86400
    kc_currentSE_mz3 = kc_init_mz3[4]
    et_currentSE_mz3 = et_init_mz3[4]
    currentStatesSE_mz3 = simulate_Richards_equation(
        x_currentSE_mz3, u_currentSE_mz3, et_currentSE_mz3, kc_currentSE_mz3, 4, rd_currentSE_mz1, lai_currentSE_mz1, soilPars_mz3)
    #rootzone_head_mz3 = compute_root_zone_moisture(currentStatesSE_mz3)
    rootzone_moist_mz3 = compute_root_zone_moisture(
        obtain_soil_moisture(currentStatesSE_mz3, soilPars_mz3), rd_currentSE_mz1)

    stateTrajectory_closedLoop_mz1.append(rootzone_moist_mz1)
    stateTrajectory_closedLoop_mz2.append(rootzone_moist_mz2)
    stateTrajectory_closedLoop_mz3.append(rootzone_moist_mz3)

    allStates_mz1[k+1, :] = currentStatesSE_mz1
    allStates_mz2[k+1, :] = currentStatesSE_mz2
    allStates_mz3[k+1, :] = currentStatesSE_mz3

    pass

#Post-processing and visualization
u_prescribed_mz1 = np.array(prescribedIrrigation_closedLoop_mz1)*1000
u_prescribed_mz2 = np.array(prescribedIrrigation_closedLoop_mz2)*1000
u_prescribed_mz3 = np.array(prescribedIrrigation_closedLoop_mz3)*1000


# Obtain the arrays for the time
timeStamps = np.arange(1, simulationPeriod+2)
timeStamps_input = np.arange(1, simulationPeriod+1)

fig, axs = plt.subplots(3, 2, figsize=(15, 15))
axs[0, 0].bar(timeStamps_input, u_prescribed_mz1, width=1.2, facecolor=None)
axs[0, 0].bar(timeStamps_input, _rain[4:len(
    u_prescribed_mz1)+4], color='red', width=1.0)
axs[0, 0].legend(['Triggered', 'Rain'])
axs[0, 0].set_title('Irrigation -  MZ 1 ')
axs[0, 0].set_ylabel('Irrigation (mm/day)')
axs[0, 0].set_xlabel('Time (days)')
# axs[0,0].grid(True,linewidth=0.01)
axs[0, 0].set_xlim(xmin=1)

axs[0, 1].plot(timeStamps, stateTrajectory_closedLoop_mz1,
               color='blue', linewidth=1.8)
axs[0, 1].plot(timeStamps, UpperZone_mz1*np.ones(len(stateTrajectory_closedLoop_mz1)),
               color='darkcyan', linestyle='dashdot', linewidth=2.5)
axs[0, 1].plot(timeStamps, LowerZone_mz1*np.ones(len(stateTrajectory_closedLoop_mz1)),
               color='darkmagenta', linestyle='dashdot', linewidth=2.5)
#axs[0,1].plot(timeStamps, np.array(LowerZones_mz1), color='darkmagenta', linestyle = 'dashdot', linewidth=2.2)
axs[0, 1].plot(timeStamps, 0.12*np.ones(len(stateTrajectory_closedLoop_mz1)),
               color='red', linestyle='dashdot', linewidth=2.5)
axs[0, 1].legend(['state Trajectory', 'FC', 'Threshold', 'PWP'])
axs[0, 1].set_title('Root Zone Soil Moisture - MZ 1')
axs[0, 1].set_ylabel('Volumetric moisture')
axs[0, 1].set_xlabel('Time (days)')
axs[0, 1].set_xlim(xmin=1)
# axs[0,1].grid(True,linewidth=0.01)

axs[1, 0].bar(timeStamps_input, u_prescribed_mz2, width=1.2)
axs[1, 0].bar(timeStamps_input, _rain[4:len(
    u_prescribed_mz1)+4], color='red', width=1.0)
axs[1, 0].legend(['Triggered', 'Rain'])
axs[1, 0].set_title('Irrigation -  MZ 2')
axs[1, 0].set_ylabel('Irrigation (mm/day)')
axs[1, 0].set_xlabel('Time (days)')
# axs[1,0].grid(True,linewidth=0.01)
axs[1, 0].set_xlim(xmin=1)

axs[1, 1].plot(timeStamps, stateTrajectory_closedLoop_mz2,
               color='blue', linewidth=1.8)
axs[1, 1].plot(timeStamps, UpperZone_mz2*np.ones(len(stateTrajectory_closedLoop_mz2)),
               color='darkcyan', linestyle='dashdot', linewidth=2.5)
axs[1, 1].plot(timeStamps, LowerZone_mz2*np.ones(len(stateTrajectory_closedLoop_mz2)),
               color='darkmagenta', linestyle='dashdot', linewidth=2.5)
#axs[1,1].plot(timeStamps, np.array(LowerZones_mz2), color='darkmagenta', linestyle = 'dashdot', linewidth=2.2)
axs[1, 1].plot(timeStamps, 0.12*np.ones(len(stateTrajectory_closedLoop_mz1)),
               color='red', linestyle='dashdot', linewidth=2.5)
axs[1, 1].legend(['state Trajectory', 'FC', 'Threshold', 'PWP'])
axs[1, 1].set_title('Root Zone Soil Moisture -  MZ 2')
axs[1, 1].set_ylabel('Volumetric moisture')
axs[1, 1].set_xlabel('Time (days)')
axs[1, 1].set_xlim(xmin=1)
# axs[1,1].grid(True,linewidth=0.01)

axs[2, 0].bar(timeStamps_input, u_prescribed_mz3, width=1.2)
axs[2, 0].bar(timeStamps_input, _rain[4:len(
    u_prescribed_mz1)+4], color='red', width=1.0)
axs[2, 0].legend(['Triggered', 'Rain'])
axs[2, 0].set_title('Irrigation - MZ 3')
axs[2, 0].set_ylabel('Irrigation (mm/day)')
axs[2, 0].set_xlabel('Time (days)')
# axs[2,0].grid(True,linewidth=0.01)
axs[2, 0].set_xlim(xmin=1)

axs[2, 1].plot(timeStamps, stateTrajectory_closedLoop_mz3,
               color='blue', linewidth=1.8)
axs[2, 1].plot(timeStamps, UpperZone_mz3*np.ones(len(stateTrajectory_closedLoop_mz3)),
               color='darkcyan', linestyle='dashdot', linewidth=2.5)
axs[2, 1].plot(timeStamps, LowerZone_mz3*np.ones(len(stateTrajectory_closedLoop_mz3)),
               color='darkmagenta', linestyle='dashdot', linewidth=2.5)
#axs[2,1].plot(timeStamps, np.array(LowerZones_mz3),color='darkmagenta', linestyle = 'dashdot', linewidth=2.2)
axs[2, 1].plot(timeStamps, 0.16*np.ones(len(stateTrajectory_closedLoop_mz1)),
               color='red', linestyle='dashdot', linewidth=2.5)
axs[2, 1].legend(['state Trajectory', 'FC', 'Threshold', 'PWP'])
axs[2, 1].set_title('Root Zone Soil Moisture -  MZ 3')
axs[2, 1].set_ylabel('Volumetric moisture')
axs[2, 1].set_xlabel('Time (days)')
# axs[2,1].grid(True,linewidth=0.01)
axs[2, 1].set_xlim(xmin=1)
fig.tight_layout()
fig.savefig("./results/triggered_mad65.pdf")

total_irrigation_mz1 = (sum(prescribedIrrigation_closedLoop_mz1))*1000
total_irrigation_mz2 = (sum(prescribedIrrigation_closedLoop_mz2))*1000
total_irrigation_mz3 = (sum(prescribedIrrigation_closedLoop_mz3))*1000

# Savings in water is
savings_mz1 = ((total_irrigation_mz1 - 80.6)/total_irrigation_mz1)*100
savings_mz2 = ((total_irrigation_mz2 - 84.8)/total_irrigation_mz2)*100
savings_mz3 = ((total_irrigation_mz3 - 94.28)/total_irrigation_mz3)*100

np.savetxt('./results/state_mz1_mad65.txt', stateTrajectory_closedLoop_mz1)
np.savetxt('./results/state_mz2_mad65.txt', stateTrajectory_closedLoop_mz2)
np.savetxt('./results/state_mz3_mad65.txt', stateTrajectory_closedLoop_mz3)

np.savetxt('./results/irrig_mz1_mad65.txt', u_prescribed_mz1)
np.savetxt('./results/irrig_mz2_mad65.txt', u_prescribed_mz2)
np.savetxt('./results/irrig_mz3_mad65.txt', u_prescribed_mz3)

# Evaluation of the cost
lower_bound_violation_mz1 = 0.176 * \
    np.ones(len(stateTrajectory_closedLoop_mz1)) - \
    stateTrajectory_closedLoop_mz1
lower_bound_violation_mz2 = 0.176 * \
    np.ones(len(stateTrajectory_closedLoop_mz2)) - \
    stateTrajectory_closedLoop_mz2
lower_bound_violation_mz3 = 0.209 * \
    np.ones(len(stateTrajectory_closedLoop_mz3)) - \
    stateTrajectory_closedLoop_mz3

upper_bound_violation_mz1 = stateTrajectory_closedLoop_mz1 - \
    0.280*np.ones(len(stateTrajectory_closedLoop_mz1))
upper_bound_violation_mz2 = stateTrajectory_closedLoop_mz2 - \
    0.280*np.ones(len(stateTrajectory_closedLoop_mz2))
upper_bound_violation_mz3 = stateTrajectory_closedLoop_mz3 - \
    0.300*np.ones(len(stateTrajectory_closedLoop_mz3))

lower_bound_violation_mz1[lower_bound_violation_mz1 < 0] = 0
lower_bound_violation_mz2[lower_bound_violation_mz2 < 0] = 0
lower_bound_violation_mz3[lower_bound_violation_mz3 < 0] = 0

upper_bound_violation_mz1[upper_bound_violation_mz1 < 0] = 0
upper_bound_violation_mz2[upper_bound_violation_mz2 < 0] = 0
upper_bound_violation_mz3[upper_bound_violation_mz3 < 0] = 0

lower_bound_violation_mz1 = lower_bound_violation_mz1**2
lower_bound_violation_mz2 = lower_bound_violation_mz2**2
lower_bound_violation_mz3 = lower_bound_violation_mz3**2

upper_bound_violation_mz1 = upper_bound_violation_mz1**2
upper_bound_violation_mz2 = upper_bound_violation_mz2**2
upper_bound_violation_mz3 = upper_bound_violation_mz3**2

cost_lower_bound_violation_mz1 = 2.0e07*sum(lower_bound_violation_mz1)
cost_lower_bound_violation_mz2 = 2.0e07*sum(lower_bound_violation_mz2)
cost_lower_bound_violation_mz3 = 2.0e07*sum(lower_bound_violation_mz3)

cost_upper_bound_violation_mz1 = 2.2e07*sum(upper_bound_violation_mz1)
cost_upper_bound_violation_mz2 = 2.2e07*sum(upper_bound_violation_mz2)
cost_upper_bound_violation_mz3 = 2.2e07*sum(upper_bound_violation_mz3)


cost_rotations_mz1 = 1000*sum(prescribedIrrigationDecision_closedLoop_mz1)
cost_rotations_mz2 = 1000*sum(prescribedIrrigationDecision_closedLoop_mz2)
cost_rotations_mz3 = 1000*sum(prescribedIrrigationDecision_closedLoop_mz3)

cost_irrigations_mz1 = 9*sum(u_prescribed_mz1)
cost_irrigations_mz2 = 9*sum(u_prescribed_mz2)
cost_irrigations_mz3 = 9*sum(u_prescribed_mz3)

total_irrigation_rate = sum(u_prescribed_mz1) + \
    sum(u_prescribed_mz2) + sum(u_prescribed_mz3)

total_irrig_cost_quadrant = cost_irrigations_mz1 + \
    cost_irrigations_mz2+cost_irrigations_mz3
total_cost_upper_bound_violation = cost_upper_bound_violation_mz1 + \
    cost_upper_bound_violation_mz2 + cost_upper_bound_violation_mz3
total_cost_lower_bound_violation = cost_lower_bound_violation_mz1 + \
    cost_lower_bound_violation_mz2 + cost_lower_bound_violation_mz3

overall_cost = 7000+total_irrig_cost_quadrant + \
    total_cost_lower_bound_violation + total_cost_upper_bound_violation

overall_cost_mz1 = cost_lower_bound_violation_mz1 + \
    cost_upper_bound_violation_mz1 + cost_rotations_mz1 + cost_irrigations_mz1
overall_cost_mz2 = cost_lower_bound_violation_mz2 + \
    cost_upper_bound_violation_mz2 + cost_rotations_mz2 + cost_irrigations_mz2
overall_cost_mz3 = cost_lower_bound_violation_mz3 + \
    cost_upper_bound_violation_mz3 + cost_rotations_mz3 + cost_irrigations_mz3
